import argparse
import os
import numpy as np
from tqdm import tqdm

def get_parser():
    """
    Set up the argument parser for the script.
    """
    parser = argparse.ArgumentParser(
        description=(
            "Corrects all camera poses in preprocessed WildRGBD data. This script "
            "assumes every frame was subjected to an incorrect global rotation "
            "instead of the correct local rotation. It undoes the wrong rotation "
            "and applies the correct one, saving the result in '_cam2.npz' files."
        )
    )
    parser.add_argument(
        "--processed_dir",
        type=str,
        default='/lc/data/3D/wildrgbd/processed',
        help="Root directory of the incorrectly processed WildRGB-D data (containing the _cam.npz files).",
    )
    return parser

def correct_pose_fully(cam_file_path):
    """
    Applies a full corrective transformation to the pose in a given camera file.
    It undoes the incorrect global rotation and applies the correct local rotation.

    Args:
        cam_file_path (str): The full path to the processed _cam.npz file.

    Returns:
        bool: True if correction was successful, False otherwise.
    """
    try:
        # --- 1. Load the processed camera data ---
        processed_data = np.load(cam_file_path)
        pose_wrong = processed_data['pose']
        intrinsics_processed = processed_data['intrinsics']

        # --- 2. Define the rotation matrices ---
        # The faulty script pre-multiplied by R_z_90_cw.
        
        # 4x4 rotation matrix for 90-degree CLOCKWISE rotation around z-axis
        R_z_90_cw = np.array([
            [ 0,  1,  0,  0],
            [-1,  0,  0,  0],
            [ 0,  0,  1,  0],
            [ 0,  0,  0,  1]
        ], dtype=np.float32)

        # The inverse of a CW rotation is a CCW rotation.
        # 4x4 rotation matrix for 90-degree COUNTER-clockwise rotation around z-axis
        R_z_90_ccw = np.array([
            [ 0, -1,  0,  0],
            [ 1,  0,  0,  0],
            [ 0,  0,  1,  0],
            [ 0,  0,  0,  1]
        ], dtype=np.float32)
        
        # The inverse matrix is R_z_90_ccw
        inv_R = R_z_90_ccw

        # --- 3. Apply the full corrective transformation ---
        # Step 1: Undo the incorrect global rotation to get the original pose.
        # pose_original = inv(R_z_90_cw) @ pose_wrong
        pose_original = inv_R @ pose_wrong

        # Step 2: Apply the correct local rotation to match the rotated image.
        # pose_final = pose_original @ R_z_90_cw
        pose_final = pose_original @ R_z_90_cw
        
        # --- 4. Save the new camera file ---
        output_path = cam_file_path.replace("_cam.npz", "_cam2.npz")
        np.savez(
            output_path,
            intrinsics=intrinsics_processed,
            pose=pose_final,
        )
        return True

    except Exception as e:
        print(f"Error processing {cam_file_path}: {e}")
        return False


def main():
    """
    Main function to find and correct all camera files.
    """
    parser = get_parser()
    args = parser.parse_args()

    processed_dir = args.processed_dir

    if not os.path.isdir(processed_dir):
        print(f"Error: Processed directory not found at '{processed_dir}'")
        return

    # Find all _cam.npz files
    cam_files_to_process = []
    for root, _, files in os.walk(processed_dir):
        for file in files:
            if file.endswith("_cam.npz"):
                cam_files_to_process.append(os.path.join(root, file))

    if not cam_files_to_process:
        print("No '_cam.npz' files found to process.")
        return

    print(f"Found {len(cam_files_to_process)} camera files to correct.")

    success_count = 0
    # Use tqdm for a progress bar
    for cam_file in tqdm(cam_files_to_process, desc="Correcting all poses"):
        if correct_pose_fully(cam_file):
            success_count += 1

    print("\nProcessing complete.")
    print(f"Total files checked: {len(cam_files_to_process)}")
    print(f"Total poses successfully corrected and saved: {success_count}")


if __name__ == "__main__":
    main()
